package at.jku.cp.mir;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

import com.google.common.collect.TreeMultimap;
import com.google.common.collect.TreeMultiset;

public class TfIdfCalculator
{

	public TfIdfCalculator()
	{
		super();
	}

	public Map<String, ArtistVector> calculateTfIdf(Map<String, List<Tag>> artistTagMap)
	{
		// this multiset contains all tags and their occurence-count
		TreeMultiset<String> tagNameMultiset = TreeMultiset.create();

		// this multimap contains all tags and their corresponding artists
		TreeMultimap<String, String> tagArtistMap = TreeMultimap.create();

		Set<String> artists = artistTagMap.keySet();
		for (String artist : artists)
		{
			List<Tag> tags = artistTagMap.get(artist);
			for (Tag tag : tags)
			{
				tagNameMultiset.add(tag.getName(), tag.getCount());
				tagArtistMap.put(tag.getName(), artist);
			}
		}

		// get the number of artists
		int D = artistTagMap.size();

		// do the TF-IDF calculation
		Set<String> tagNames = tagNameMultiset.elementSet();
		Map<String, ArtistVector> artistVectors = new TreeMap<String, ArtistVector>();
		for (String tagName : tagNames)
		{
			// get all artists which have this tag
			Set<String> artistsForThisTag = tagArtistMap.get(tagName);
			double nrOfArtistsWithThisTag = artistsForThisTag.size();

			for (String artist : artistsForThisTag)
			{
				int n = getTagCount(artistTagMap, artist, tagName);
				double nrOfTagsForThisArtist = artistTagMap.get(artist).size();
				double tf = n / nrOfTagsForThisArtist;
				double idf = Math.log(D / (1 + nrOfArtistsWithThisTag));

				double tfidf = tf * idf;

				ArtistVector aV = artistVectors.get(artist);
				if (aV == null)
				{
					aV = new ArtistVector(artist);
					artistVectors.put(artist, aV);
				}
				aV.addTagTfIdf(tagName, tfidf);
			}
		}

		return artistVectors;
	}

	private int getTagCount(Map<String, List<Tag>> artistTagMap, String artist, String tagName)
	{
		List<Tag> list = artistTagMap.get(artist);
		for (Tag tag : list)
		{
			if (tag.getName().equalsIgnoreCase(tagName))
				return tag.getCount();
		}

		return 0;
	}

}
